package com.ht.modules.component.ai.chat.service;

import com.ht.modules.component.ai.assistants.Assistant;
import com.ht.modules.component.ai.assistants.IAssistantService;
import com.ht.modules.component.ai.chat.entity.ChatRequestParams;
import com.ht.modules.component.ai.session.assistant.SessionChatAssistant;
import com.ht.modules.component.ai.session.entity.SessionResponse;
import com.ht.modules.component.ai.session.service.SessionDbService;
import com.ht.modules.component.ai.tools.service.ToolService;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.service.TokenStream;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

import static com.ht.modules.component.ai.session.entity.SessionResponse.ResponseStage.*;
import static com.ht.modules.component.ai.session.entity.SessionResponse.THINK_END_TAG;
import static com.ht.modules.component.ai.session.entity.SessionResponse.THINK_START_TAG;

@Slf4j
@Service
public class ChatService {
    @Autowired
    private SessionDbService sessionDbService;
    @Autowired
    private StreamingChatModel streamingChatModel;
    @Autowired
    private ToolService toolService;
    @Autowired
    private IAssistantService<Assistant> assistantService;

    /**
     * 会话管理
     */
    private final Map<String, SessionChatAssistant> chatOptionsMap = new ConcurrentHashMap<>(16);

    /**
     * 重置聊天
     *
     * @param chatRequestParams 参数
     * @return {@link String}
     */
    public String resetChat(ChatRequestParams chatRequestParams) {
        chatOptionsMap.remove(chatRequestParams.getId());
        sessionDbService.deleteMessages(chatRequestParams.getId());
        return "操作成功";
    }


    /**
     * 流式聊天
     *
     * @param chatRequestParams 参数
     * @return 流式响应
     */
    public Flux<?> chatFlux(ChatRequestParams chatRequestParams) {
        try {
            String message = chatRequestParams.getMessage();
            TokenStream tokenStream = this.getAssistant(chatRequestParams).chatStream(chatRequestParams.getId(), message);
            return this.fluxConversion(tokenStream).delayElements(Duration.ofMillis(40));
        } catch (Exception e) {
            return Flux.error(e);
        }
    }

    /**
     * 初始化-会话监听
     */
    @PostConstruct
    private void init() {
        new Thread(() -> {
            while (true) {
                try {
                    TimeUnit.SECONDS.sleep(3);
                    StringBuilder builder = new StringBuilder();
                    for (String key : chatOptionsMap.keySet()) {
                        SessionChatAssistant sessionChatAssistant = chatOptionsMap.get(key);
                        long createTime = sessionChatAssistant.getTime();
                        if (System.currentTimeMillis() - createTime > 120 * 1000) {
                            chatOptionsMap.remove(key);
                            continue;
                        }
                        builder.append("\t").append(key);
                    }
                    int size = chatOptionsMap.size();
                    log.info("当前会话数量:{};会话ID{}", size, builder);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
            }
        }, "===会话监听===").start();
    }

    /**
     * 流式响应转换
     *
     * @param tokenStream 令牌流
     * @return {@link Flux}<{@link SessionResponse}> 流式响应
     */
    private Flux<SessionResponse> fluxConversion(TokenStream tokenStream) {
        AtomicReference<SessionResponse.ResponseStage> stage = new AtomicReference<>(INITIAL);
        AtomicReference<SessionResponse.ResponseStage> flag = new AtomicReference<>(INITIAL);
        AtomicReference<SessionResponse.ResponseStage> think = new AtomicReference<>(INITIAL);
        return Flux.create(sink -> tokenStream
                .onPartialResponse(partialResponse -> {
                    if (partialResponse.contains(THINK_START_TAG)) {
                        stage.set(THINKING);
                    }
                    if (partialResponse.contains(THINK_END_TAG)) {
                        stage.set(POST_THINKING);
                    }
                    SessionResponse sessionResponse;
                    switch (stage.get()) {
                        case THINKING -> {
                            if (think.get() == INITIAL) {
                                partialResponse = partialResponse.replaceAll("<think>", "<思考>");
                            } else {
                                partialResponse = partialResponse.replaceAll("<think>", "<深度思考>");
                            }
                            sessionResponse = SessionResponse.think(partialResponse);
                            sink.next(sessionResponse);
                        }
                        case POST_THINKING -> {
                            if (think.get() == INITIAL) {
                                partialResponse = partialResponse.replaceAll("</think>", "</思考>");
                            } else {
                                partialResponse = partialResponse.replaceAll("</think>", "</深度思考>");
                            }
                            think.set(POST_THINKING);
                            stage.set(REPLYING);
                            flag.set(POST_THINKING);
                            sessionResponse = SessionResponse.think(partialResponse);
                            sink.next(sessionResponse);
                            sink.next(SessionResponse.think("\n\n"));
                        }
                        default -> {
                            if (flag.get() == POST_THINKING) {
                                partialResponse = partialResponse.replaceAll("\n\n", "");
                                flag.set(INITIAL);
                            }
                            sessionResponse = SessionResponse.reply(partialResponse);
                            sink.next(sessionResponse);
                        }
                    }
                })
                .onRetrieved(r -> {
                    log.info("通过向量数据库检索到的结果:::");
                    for (Content content : r) {
                        log.info("\n{}", content);
                    }
                })
                .onToolExecuted(toolExecuted -> {
                    String result = toolExecuted.result();
                    System.err.println(result);
                })
                .onCompleteResponse(completeResponse -> {
                    sink.next(SessionResponse.end());
                    sink.complete();
                })
                .onError(error -> {
                    log.error("error = {}", String.valueOf(error));
                    sink.next(SessionResponse.error(error.getMessage()));
                    sink.complete();
                })
                .start());
    }

    /**
     * 获取助手
     *
     * @param chatRequestParams 参数
     * @return {@link Assistant}
     */
    private Assistant getAssistant(ChatRequestParams chatRequestParams) {
        // 获取会话
        SessionChatAssistant sessionChatAssistant = null;
        if (chatOptionsMap.containsKey(chatRequestParams.getId())) {
            sessionChatAssistant = chatOptionsMap.get(chatRequestParams.getId());
            sessionChatAssistant.setTime(System.currentTimeMillis());
        }
        chatRequestParams.setMessage("");
        if (null == sessionChatAssistant || !chatRequestParams.equals(sessionChatAssistant.getChatRequestParams())) {
            // TODO 这里到时候可以通过会话信息去确定需要从哪个账号的知识库去检索
            sessionChatAssistant = SessionChatAssistant.builder()
                    .time(System.currentTimeMillis())
                    .assistant(assistantService.buildAssistant(
                            streamingChatModel,
                            assistantService.buildChatMemoryProvider(5),
                            assistantService.buildEmbeddingStoreContentRetriever("disease", 3, 0.7, 0.8, 1),
                            toolService,
                            assistantService.buildToolProvider("McpServers.json"),
                            Assistant.class)
                    )
                    .build();
            sessionChatAssistant.setChatRequestParams(chatRequestParams);
            chatOptionsMap.put(chatRequestParams.getId(), sessionChatAssistant);
        }
        return sessionChatAssistant.getAssistant();
    }
}
